#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
模板数据导入脚本
项目：南华大学附属第一医院康复医学科管理系统
创建日期：2025-11-21
说明：此脚本用于导入初始模板数据到 templates 表
"""

import psycopg2
import sys
import os

# 数据库连接配置
DB_CONFIG = {
    'host': 'localhost',#改成服务器IP
    'port': 5432,
    'database': 'nanhua_rehab',
    'user': 'postgres',
    'password': '123456'  # 请修改为实际密码
}


def insert_templates():
    """插入模板数据到数据库"""
    conn = None
    cursor = None
    
    try:
        print("=" * 60)
        print("开始导入模板数据...")
        print("=" * 60)
        
        # 连接数据库
        print(f"\n>> 正在连接数据库 {DB_CONFIG['database']}...")
        conn = psycopg2.connect(**DB_CONFIG)
        conn.autocommit = False
        cursor = conn.cursor()
        
        print(">> 数据库连接成功！\n")
        
        # 检查 templates 表是否存在
        cursor.execute("""
            SELECT EXISTS (
                SELECT FROM information_schema.tables 
                WHERE table_schema = 'public' 
                AND table_name = 'templates'
            );
        """)
        
        if not cursor.fetchone()[0]:
            print("错误：templates 表不存在！")
            print("   请先运行 01_create_tables.py 创建表结构")
            sys.exit(1)
        
        # 清空现有数据（可选）
        print("清空现有模板数据...")
        cursor.execute("TRUNCATE TABLE templates RESTART IDENTITY CASCADE;")
        print(" 清空完成\n")
        
        # 获取脚本所在目录
        script_dir = os.path.dirname(os.path.abspath(__file__))
        template_file = os.path.join(script_dir, 'pgsql_localhost_-2025_11_21_00_46_00-dump.sql')
        
        print(">>> 正在读取模板数据文件...")
        
        try:
            # 使用 PostgreSQL 的 COPY 命令直接导入
            # 这样可以正确处理转义字符
            with open(template_file, 'r', encoding='utf-8') as f:
                content = f.read()
            
            # 找到 COPY 命令的开始和结束
            copy_start = content.find('COPY public.templates')
            if copy_start == -1:
                print(" 错误：文件格式不正确，找不到 COPY 命令")
                sys.exit(1)
            
            # 提取 COPY 数据部分（从 FROM stdin; 之后到 \. 之前）
            data_start = content.find('FROM stdin;', copy_start) + len('FROM stdin;')
            data_end = content.find('\n\\.', data_start)
            
            if data_end == -1:
                data_end = len(content)
            
            # 计算有多少行数据
            data_section = content[data_start:data_end].strip()
            data_lines = [line for line in data_section.split('\n') if line.strip()]
            
            print(f"OK 找到 {len(data_lines)} 条模板记录\n")
            print(" 开始使用 COPY 命令导入数据...")
            
            # 使用 COPY 命令导入（最可靠的方式）
            copy_sql = """
                COPY templates (id, name, description, category, latest_version, is_active, 
                               created_by, created_at, updated_at, version_history)
                FROM STDIN WITH (FORMAT text, DELIMITER E'\\t', ENCODING 'UTF8');
            """
            
            # 准备数据
            from io import StringIO
            data_buffer = StringIO(data_section)
            
            cursor.copy_expert(copy_sql, data_buffer)
            
            # 提交事务
            conn.commit()
            
            print(f"OK 成功导入数据！\n")
            
        except FileNotFoundError:
            print(f" 错误：找不到模板数据文件 {template_file}")
            print("   请确保文件位于正确的路径")
            sys.exit(1)
        except Exception as e:
            print(f" 导入失败：{e}")
            import traceback
            traceback.print_exc()
            sys.exit(1)
        
        # 验证插入结果
        print(" 验证插入结果...")
        cursor.execute("SELECT id, name FROM templates ORDER BY id;")
        templates = cursor.fetchall()
        
        print(f"\n 已导入的模板（共 {len(templates)} 个）：")
        for template_id, name in templates:
            print(f"   {template_id}. {name}")
        
        print("\n" + "=" * 60)
        print(" 模板数据导入完成！")
        print("=" * 60)
        
    except psycopg2.Error as e:
        if conn:
            conn.rollback()
        print(f"\n 数据库错误：{e}")
        print(f"   错误代码：{e.pgcode}")
        print(f"   错误详情：{e.pgerror}")
        sys.exit(1)
        
    except Exception as e:
        if conn:
            conn.rollback()
        print(f"\n 发生错误：{e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)
        
    finally:
        # 关闭连接
        if cursor:
            cursor.close()
        if conn:
            conn.close()
            print("\n🔌 数据库连接已关闭")


if __name__ == "__main__":
    print("""
╔════════════════════════════════════════════════════════════╗
║   南华大学附属第一医院康复医学科管理系统                   ║
║   模板数据导入脚本                                         ║
╚════════════════════════════════════════════════════════════╝
    """)
    
    # 检查配置
    if DB_CONFIG['password'] == 'your_password_here':
        print("!!  警告：请先在脚本中修改数据库密码！")
        print("   编辑 DB_CONFIG['password'] 字段\n")
        sys.exit(1)
    
    # 确认执行
    print(f" 数据库配置：")
    print(f"   主机：{DB_CONFIG['host']}:{DB_CONFIG['port']}")
    print(f"   数据库：{DB_CONFIG['database']}")
    print(f"   用户：{DB_CONFIG['user']}")
    print()
    
    response = input("!!  此操作将清空现有模板数据并重新导入！是否继续？(yes/no): ")
    
    if response.lower() in ['yes', 'y']:
        insert_templates()
    else:
        print("\n 操作已取消")
        sys.exit(0)
